﻿#pragma once

#include <iostream>
#include <vector>
#include <fstream>
#include <chrono>
#include <numeric>
#include <algorithm>
#include <filesystem>
#include "wavreader.h"
#include "svm.h"
#include "svm_train.h"
#include "svm_predict.h"
#include "tools.h"
#include "wavAudio.h"

using namespace std;

#define USE_500_DATASET true
#define GAMMA 0.056
#define C 1.198

int result = -1; // 储存预测结果
bool is_train = false;
bool is_test_accuracy = false;
float accuracy = 0;

void writeToTxt(const std::string& filepath, int label, const std::vector<float>& character) {
	std::ofstream outfile(filepath, std::ios::app);
	if (outfile.is_open()) {
		outfile << label << " ";
		int character_size = character.size();
		//int step = character_size / length;
		for (int i = 0; i < character_size; i++) {
			outfile << i << ":" << character[i];
			if (i != (character_size - 1)) {
				outfile << " ";
			}
		}
		outfile << std::endl;
		outfile.close();
	}
	else {
		std::cerr << "Failed to open file." << std::endl;
	}
}

void computeMelspectrogram(const std::string& folderPath, int numFiles, int label) {
	// 构造文件路径并计算梅尔频谱图
	std::cout << "\r" << std::string(80, ' ') << "\r";
	for (int i = 0; i < numFiles; ++i) {
		if (i % 5 == 0 || i == numFiles - 1) {
			displayProgressBar(i + 1, numFiles);
		}
		std::string filePath = folderPath + std::to_string(i) + ".wav";
		//std::cout << filePath << std::endl;
		//std::vector<float> melspectrogram = computeMelspectrogram(filePath);

		// 将数据降维到指定的维度
		//std::vector<float> reduced_data = pca_reduce_dimension(melspectrogram, SV_DIM);

		// 将梅尔频谱图的长度存储到特征文件中
		//writeToTxt("data/character.txt", label, reduced_data);
		wavAudio obj(filePath.c_str());
		obj.computeCharacter();
		std::vector<float> tmp = obj.returnCharacter();
		if (USE_500_DATASET) {
			writeToTxt("data/character-500.txt", label, tmp);
		}
		else {
			writeToTxt("data/character.txt", label, tmp);
		}
	}
}


void computeCharacter() {
	std::string folderPath0;
	std::string folderPath1;
	int numFiles;
	if (USE_500_DATASET) {
		folderPath0 = "data/0-500/0_";
		folderPath1 = "data/1-500/1_";
		numFiles = 500;
	}
	else {
		folderPath0 = "data/0/0_";
		folderPath1 = "data/1/1_";
		numFiles = 6000;
	}

	// 计算data/0/文件夹下文件的梅尔频谱图长度
	std::cout << "Compute Mel Spectrogram from " << folderPath0 << std::endl;
	computeMelspectrogram(folderPath0, numFiles, 0);

	// 计算data/1/文件夹下文件的梅尔频谱图长度
	std::cout << std::endl << "Compute Mel Spectrogram from " << folderPath1 << std::endl;
	computeMelspectrogram(folderPath1, numFiles, 1);
	std::cout << std::endl;
}

//训练模型
void trainModel()
{
	is_train = true;
	std::cout << "********************模型训练********************" << endl;
	clock_t start = clock(); //begin time
	while (true)
	{
		std::cout << "是否需要计算特征（Y or N)";
		char a;
		std::cin >> a;
		if (a == 'y' || a == 'Y') {
			std::cout << "* 【开始计算特征】                             *" << endl;
			if (USE_500_DATASET) {
				std::remove("data/character-500.txt");
			}
			else {
				std::remove("data/character.txt");
			}
			computeCharacter();
			std::cout << endl << "* 【特征计算完毕】                             *" << endl;
			break;
		}
		else if (a == 'n' || a == 'N') {
			break;
		}
		else {
			std::cout << "输入错误" << std::endl;
		}
	}
	int paramNumberPredict = 11;
	char* tmp[11] = { "" };
	tmp[1] = { "-s" };
	tmp[2] = { "0" };//"c_svc","nu_svc","one_class","epsilon_svr","nu_svr",NULL
	tmp[3] = { "-t" };
	tmp[4] = { "2" };//"linear","polynomial","rbf","sigmoid","precomputed",NULL
	tmp[5] = { "-c" }; //C
	tmp[6] = new char[10]; // 假设文件路径不会超过 10 个字符
	sprintf(tmp[6], "%f", C);
	tmp[7] = { "-g" }; //gamma
	tmp[8] = new char[10]; // 假设文件路径不会超过 10 个字符
	sprintf(tmp[8], "%f", GAMMA);
	if (USE_500_DATASET) {
		tmp[9] = { "data/character-500.txt" };
		tmp[10] = { "data/SnoringModel-500.txt" };
	}
	else {
		tmp[9] = { "data/character.txt" };
		tmp[10] = { "data/SnoringModel.txt" };
	}
	std::cout << "* 【开始训练模型，请稍后】                   *" << endl;
	svmTrain(paramNumberPredict, tmp);
	std::cout << "* 【模型训练成功,特征数目:" << 50 << " 】                        *" << endl;

	//std::cout << "* 【未成功识别笔画数的数据数目】  " << numOfNotSplit << "           *" << endl;
	std::cout << "* 【指令集类别总数】：2 (0: 非鼾声, 1:鼾声)  *" << endl;
	//printTrainCount();
	clock_t end = clock(); // end time
	std::cout << "训练耗时：" << double(end - start) / CLOCKS_PER_SEC << "s" << endl;  //输出时间（单位：ｓ）
	std::cout << "*********************end************************" << endl << endl;
	is_train = false;
}

void predict() {
	string path;
	while (true)
	{
		std::cout << "请输入文件路径(eg:data/0/0_1.wav)";
		std::cin >> path;
		if (fileExists(path)) {
			result = -1;
			wavAudio obj(path.c_str());
			obj.computeCharacter();
			std::vector<float> character = obj.returnCharacter();
			//for (int i = 0; i < character.size(); i++) {
			//	std::cout << character[i] << " ";
			//}
			//std::cout << std::endl;
			std::remove("data/singleCharacter.txt");
			writeToTxt("data/singleCharacter.txt", 0, character);

			// 开始预测
			int paramNumberPredict = 4;
			char* tmp[4] = { "" };
			tmp[1] = { "data/singleCharacter.txt" };
			if (USE_500_DATASET) {
				tmp[2] = { "data/SnoringModel.txt" };
			} else{
				tmp[2] = { "data/SnoringModel-500.txt" };
			}
			cout << "* 【调用模型文件为" << tmp[2] << "】 *" << endl;
			tmp[3] = { "data/result.txt" };
			svmPredict(paramNumberPredict, tmp);
			if (result == 1) {
				std::cout << "识别到鼾声" << std::endl;
			}
			else if (result == 0) {
				std::cout << "未识别到鼾声" << std::endl;
			}
			else {
				std::cout << "未识别" << std::endl;
			}
			return;
		}
		else {
			std::cout << "文件不存在" << std::endl;
		}
	}
}

int testPredict(const std::string& filepath) {

	if (fileExists(filepath)) {
		result = -1;
		wavAudio obj(filepath.c_str());
		obj.computeCharacter();
		std::vector<float> character = obj.returnCharacter();
		//for (int i = 0; i < character.size(); i++) {
		//	std::cout << character[i] << " ";
		//}
		//std::cout << std::endl;
		std::remove("data/singleCharacter.txt");
		writeToTxt("data/singleCharacter.txt", 0, character);

		// 开始预测
		int paramNumberPredict = 4;
		char* tmp[4] = { "" };
		tmp[1] = { "data/singleCharacter.txt" };
		if (USE_500_DATASET) {
			tmp[2] = { "data/SnoringModel-500.txt" };
		}
		else {
			tmp[2] = { "data/SnoringModel.txt" };
		}
		if (!is_test_accuracy) {
			cout << "* 【调用模型文件为" << tmp[2] << "】 *" << endl;
		}
		tmp[3] = { "data/result.txt" };
		svmPredict(paramNumberPredict, tmp);
		if (result == 1) {
			if (!is_test_accuracy) {
				std::cout << "识别到鼾声" << std::endl;
			}
			return 1;
		}
		else if (result == 0) {
			if (!is_test_accuracy) {
				std::cout << "未识别到鼾声" << std::endl;
			}
			return 0;
		}
		else {
			if (!is_test_accuracy) {
				std::cout << "未识别" << std::endl;
			}
			return -1;
		}

	}
	else {
		if (!is_test_accuracy) {
			std::cout << "文件不存在" << std::endl;
		}
	}
	return -1;
}

void testAccuracy() {
	is_test_accuracy = true;
	std::string folderPath0;
	std::string folderPath1;
	int num = 0;
	char chr;
	std::cout << "请输入要验证的测试集（0：500  1：6000）：";
	cin >> chr;
	if (chr == '0') {
		folderPath0 = "data/0-500/0_";
		folderPath1 = "data/1-500/1_";
		num = 500;
	}
	else {
		folderPath0 = "data/0/0_";
		folderPath1 = "data/1/1_";
		num = 6000;
	}
	std::cout << "使用" << num << "数据集进行验证" << std::endl;
	

	int right = 0;
	int wrong = 0;
	int none = 0;

	std::remove("data/testResult.txt");
	std::ofstream outfile("data/testResult.txt", std::ios::app);
	// 计算data/0/文件夹下文件的预测正确率
	std::cout << "begin predicting the files under folder " << folderPath0 << std::endl;
	std::cout << "\r" << std::string(80, ' ') << "\r";
	for (int i = 0; i < num; i++) {
		if (i % 5 == 0 || i == num - 1) {
			displayProgressBar(i + 1, num);
		}
		string path = folderPath0 + std::to_string(i) + ".wav";
		int a = testPredict(path);
		outfile << a << std::endl;
		if (a == 0) {
			right++;
		}
		else if (a == 1) {
			wrong++;
		}
		else {
			none++;
		}
	}

	// 计算data/1/文件夹下文件的预测正确率
	std::cout << std::endl << "begin predicting the files under folder " << folderPath1 << std::endl;
	std::cout << "\r" << std::string(80, ' ') << "\r";
	for (int i = 0; i < num; i++) {
		if (i % 5 == 0 || i == num - 1) {
			displayProgressBar(i + 1, num);
		}
		string path = folderPath1 + std::to_string(i) + ".wav";
		int a = testPredict(path);
		outfile << a << std::endl;
		if (a == 1) {
			right++;
		}
		else if (a == 0) {
			wrong++;
		}
		else {
			none++;
		}
	}

	outfile.close();
	std::cout << std::endl;
	std::cout << "right:" << right << "     wrong:" << wrong << "     none:" << none << std::endl;
	std::cout << "Accuracy:" << float(float(right) / float(right + wrong + none)) * 100 << "%" << std::endl;
	is_test_accuracy = false;
}

void crossValidation() {
	clock_t start = clock(); //begin time
	int paramNumberPredict = 13;
	char* tmp[13] = { "" };
	tmp[1] = { "-s" };
	//tmp[2] = { "0" };//"c_svc","nu_svc","one_class","epsilon_svr","nu_svr",NULL
	tmp[2] = { "0" };
	tmp[3] = { "-t" };
	tmp[4] = { "2" };//"linear","polynomial","rbf","sigmoid","precomputed",NULL
	tmp[5] = { "-c" }; //C
	tmp[6] = new char[10]; // 假设文件路径不会超过 10 个字符
	sprintf(tmp[6], "%f", C);
	tmp[7] = { "-g" }; //gamma
	tmp[8] = new char[10]; // 假设文件路径不会超过 10 个字符
	sprintf(tmp[8], "%f", GAMMA);
	tmp[9] = { "-v" }; //cross validation
	tmp[10] = { "5" };
	if (USE_500_DATASET) {
		tmp[11] = { "data/character-500.txt" };
		tmp[12] = { "data/SnoringModel-500.txt" };
	}
	else {
		tmp[11] = { "data/character.txt" };
		tmp[12] = { "data/SnoringModel.txt" };
	}
	svmTrain(paramNumberPredict, tmp);
	clock_t end = clock(); // end time
	std::cout << "Test time：" << double(end - start) / CLOCKS_PER_SEC << "s" << endl;  //输出时间（单位：ｓ）
	std::cout << "Accuracy:" << accuracy << "%" << std::endl;
}

void readDataFromWAV(const char* filepath) {
	std::cout << "begin to read..." << std::endl;
	void* obj = wav_read_open(filepath);

	int format, channels, sr, bits_per_sample;
	unsigned int data_length;
	int res = wav_get_header(obj, &format, &channels, &sr, &bits_per_sample, &data_length);
	if (!res)
	{
		cerr << "get ref header error: " << res << endl;
		return;
	}

	int samples = data_length * 8 / bits_per_sample;
	std::vector<int16_t> tmp(samples);
	res = wav_read_data(obj, reinterpret_cast<unsigned char*>(tmp.data()), data_length);
	if (res < 0)
	{
		std::cerr << "read wav file error: " << res << std::endl;
		return;
	}

	std::remove("data/wavInfo.txt");
	std::ofstream outfile("data/wavInfo.txt");
	outfile << "file path:" << filepath << std::endl;
	outfile << "data length:" << data_length << std::endl;
	outfile << "audio data format:" << format << std::endl;
	outfile << "channels number:" << channels << std::endl;
	outfile << "sample rate:" << sr << std::endl;
	outfile << "bits per sample:" << bits_per_sample << std::endl;
	outfile << "samples number:" << samples << std::endl << std::endl;

	std::cout << "\r" << std::string(80, ' ') << "\r";
	for (int i = 0; i < samples; i++) {
		if (i % 10000 == 0 || i == samples - 1) {
			displayProgressBar(i + 1, samples);
		}
		outfile << tmp[i] << " ";
	}
	outfile << std::endl;
	outfile.close();
	std::cout << std::endl << "read info done." << std::endl;
	wav_read_close(obj);
}

// 判断文件中是否存在相同参数组合的函数
bool isParamExist(const std::string& filename, float g, float c) {
	std::ifstream file(filename);
	if (!file.is_open()) {
		return false;
	}

	std::string line;
	while (std::getline(file, line)) {
		std::istringstream lineStream(line);
		std::string cell;

		// 读取第一个单元格
		if (!std::getline(lineStream, cell, ',')) {
			continue; // 如果读取失败，跳过该行
		}
		float param1;
		try {
			param1 = std::stof(cell);
		}
		catch (const std::invalid_argument& e) {
			continue; // 如果转换失败，跳过该行
		}

		// 读取第二个单元格
		if (!std::getline(lineStream, cell, ',')) {
			continue; // 如果读取失败，跳过该行
		}
		float param2;
		try {
			param2 = std::stof(cell);
		}
		catch (const std::invalid_argument& e) {
			continue; // 如果转换失败，跳过该行
		}

		// 检查参数是否匹配
		if (param1 == g && param2 == c) {
			file.close();
			return true;
		}
	}

	file.close();
	return false;
}

// 写入一行数据到CSV文件
void writeRowToCSV(const std::string& filename, float g, float c, float accuracy) {
	std::ofstream file(filename, std::ios_base::app); // 打开文件追加写入模式
	if (!file.is_open()) {
		std::cerr << "Error: Unable to open file for writing." << std::endl;
		return;
	}

	// 将数据写入文件
	file << g << "," << c << "," << accuracy << std::endl;

	file.close();
}

void findpara() {
	string filename;
	if (USE_500_DATASET) {
		filename = "data/GandC-500.csv";
	}
	else {
		filename = "data/GandC.csv";
	} 
	for (float g = 0.05; g <= 0.07; g += 0.001) {
		//float g = 0.06;
		for (float c = 1.18; c <= 1.20; c += 0.001) {
			//float c = 0.054;
			// 如果参数组合已经存在于文件中，则跳过
			if (isParamExist(filename, g, c)) {
				continue;
			}

			accuracy = 0;
			int paramNumberPredict = 13;
			char* tmp[13] = { "" };
			tmp[1] = { "-s" };
			//tmp[2] = { "0" };//"c_svc","nu_svc","one_class","epsilon_svr","nu_svr",NULL
			tmp[2] = { "0" };
			tmp[3] = { "-t" };
			tmp[4] = { "2" };//"linear","polynomial","rbf","sigmoid","precomputed",NULL
			tmp[5] = { "-c" }; //C
			tmp[6] = new char[10]; // 假设文件路径不会超过 10 个字符
			sprintf(tmp[6], "%f", c);
			tmp[7] = { "-g" }; //gamma
			tmp[8] = new char[10]; // 假设文件路径不会超过 10 个字符
			sprintf(tmp[8], "%f", g);
			tmp[9] = { "-v" }; //cross validation
			tmp[10] = { "5" };
			if (USE_500_DATASET) {
				tmp[11] = { "data/character-500.txt" };
				tmp[12] = { "data/SnoringModel-500.txt" };
			}
			else {
				tmp[11] = { "data/character.txt" };
				tmp[12] = { "data/SnoringModel.txt" };
			}
			svmTrain(paramNumberPredict, tmp);
			// 写入一行数据到CSV文件
			writeRowToCSV(filename, g, c, accuracy);
		}
	}
}

int main() {
	while (1) {
		std::cout << "请输入选择：" << std::endl;
		std::cout << "0(退出)" << std::endl;
		std::cout << "1(训练模型）" << std::endl;
		std::cout << "2(单次测试）" << std::endl;
		std::cout << "3(交叉验证)" << std::endl;
		std::cout << "4(读取音频文件信息)" << std::endl;
		std::cout << "5(测试准确率)" << std::endl;
		std::cout << "6(Test Function)" << std::endl;
		std::cout << "请输入：";
		int ans;
		std::cin >> ans;
		char path[100] = "";
		switch (ans)
		{
		case 0:
			return 0; // 退出程序
		case 1:
			trainModel(); // 调用训练模型函数
			break;
		case 2:
			// TODO: 添加测试功能的代码
			predict();
			break;
		case 3:
			crossValidation();
			break;
		case 4:
			std::cout << "请输入文件路径（eg:data/0/0_0.wav）";
			std::cin >> path;
			readDataFromWAV(path);
			break;
		case 5:
			testAccuracy();
			break;
		case 6:
			findpara();
			break;
		default:
			std::cout << "输入无效，请重新输入" << std::endl;
			break;
		}
	}
	// 此处的代码将永远不会被执行，因为在 case 0 处已经退出了程序
	system("pause"); // 不需要在这里暂停
	return 0;
}